import os
import logging
import importlib
from typing import Dict
from .data_utils import TextHelper,jload,jdump
from .dalchemy_data import DalchemyData
from langchain.document_loaders import TextLoader, DirectoryLoader, CSVLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

DATA_REGISTRY: Dict[str, DalchemyData] = {}
'''
data里多写些loader，然后由dalchemy来调用，构造prompt
'''

def register_data(name):
    def register_data_fn(cls):
        assert issubclass(cls, DalchemyData)
        if cls in DATA_REGISTRY.values():
            raise ValueError("Cannot register duplicate data ({})".format(name))
        DATA_REGISTRY[name] = cls
        return cls

    return register_data_fn


def import_models(infer_dir, namespace):
    for file in os.listdir(infer_dir):
        path = os.path.join(infer_dir, file)
        if (
                not file.startswith("_")
                and not file.startswith(".")
                and (file.endswith(".py") or os.path.isdir(path))
        ):
            infer_name = file[: file.find(".py")] if file.endswith(".py") else file
            importlib.import_module(namespace + "." + infer_name)


# automatically import any Python files in the infer/ directory
infer_dir = os.path.dirname(__file__)
import_models(infer_dir, "dalchemy.data")


def build_data(name):
    return DATA_REGISTRY.get(name, DalchemyData)()
